from contextlib import asynccontextmanager, contextmanager
from dataclasses import dataclass
import pathlib
from typing import *
import logging
import trio
import trio_cdp
from cdp import debugger, runtime

import rich
from rich import box
from rich.syntax import Syntax
from rich.panel import Panel
from rich.layout import Layout
from rich.console import Console, RenderableType
from rich.logging import RichHandler
from rich.table import Table, Column
from rich.text import TextType, Text
from rich.pretty import Pretty
from rich.style import Style

T = TypeVar("T")
TEvent = TypeVar("TEvent")
T_JSON_DICT = Dict[str, Any]

ATTEMPT_COUNT=60

__import__("fixer")  # CallFrame not containts `this`
STDOUT = rich.get_console()
STDERR = Console(stderr=True)


logging.basicConfig(
    level=logging.INFO,
    format="%(message)s",
    datefmt="[%X]",
    handlers=[
        RichHandler(show_path=False, show_time=True, console=STDERR, markup=True)
    ],
)
LOG: logging.Logger = logging.getLogger("ark-debugger")


class ArkConnectError(Exception):
    error: OSError
    pass


class ArkCdp:

    def __init__(self, conn: trio_cdp.CdpConnection, nursery: trio.Nursery) -> None:
        self._conn = conn
        self._nursery = nursery
        self.context_id: runtime.ExecutionContextId

    @property
    def connection(self) -> trio_cdp.CdpConnection:
        return self._conn

    def listen(self, t: T, handler: Callable[[T], None]):
        async def a():
            async for event in self._conn.listen(t):
                await handler(event)

        self._nursery.start_soon(a)

    async def execute(self, cmd: Generator[dict, T, Any]) -> T:
        return await self._conn.execute(cmd)

    @asynccontextmanager
    async def wait_for(
        self,
        event_type: Type[T],
        buffer_size=10,
    ) -> AsyncGenerator[trio_cdp.CmEventProxy, None]:
        async with self._conn.wait_for(event_type, buffer_size) as proxy:
            yield proxy

    async def execute(self, cmd: Generator[dict, T, Any]) -> T:
        return await self._conn.execute(cmd)


@asynccontextmanager
async def open_arkcdp(url) -> AsyncIterator[ArkCdp]:
    async with trio.open_nursery() as nursery:
        conn = None
        counter = ATTEMPT_COUNT
        while conn is None:
            counter -= 1
            try:
                conn = await trio_cdp.connect_cdp(nursery, url)
            except OSError as e:
                if counter:
                    LOG.warn("%s %s",e, f"[{ATTEMPT_COUNT-counter}/{ATTEMPT_COUNT}]")
                    await trio.sleep(1)
                else:
                    raise ArkConnectError(str(e))
        context = ArkCdp(conn, nursery)
        try:
            with trio_cdp.connection_context(conn):
                yield context
        finally:
            await conn.aclose()
            nursery.cancel_scope.cancel()


class ArkDebugger:
    arkcdp: ArkCdp
    console: rich.console.Console
    execution_context: runtime.ExecutionContextDescription
    scripts: Dict[runtime.ScriptId, debugger.ScriptParsed]
    script_sources: Dict[runtime.ScriptId, str]
    paused: debugger.Paused
    debuggerId: runtime.UniqueDebuggerId

    def __init__(self, console: Console = None) -> None:
        self.console = console if console is not None else rich.get_console()
        self.scripts = dict()
        self.script_sources = dict()
        pass

    @asynccontextmanager
    async def connect(self, uri: str) -> AsyncIterator[ArkCdp]:
        async with open_arkcdp(uri) as arkcdp:
            self._listen(arkcdp)
            await self.execute(runtime.enable())
            self.debuggerId = await self.arkcdp.execute(debugger.enable(10000000))
            await self.execute(debugger.set_pause_on_exceptions("none"))
            await self.execute_and_wait_paused(runtime.run_if_waiting_for_debugger())
            yield arkcdp

    @asynccontextmanager
    async def wait_for(
        self,
        event_type: Type[T],
        buffer_size=10,
    ) -> AsyncGenerator[trio_cdp.CmEventProxy, None]:
        async with self.arkcdp.wait_for(event_type, buffer_size) as proxy:
            yield proxy

    async def execute(self, cmd: Generator[dict, T, Any]) -> T:
        return await self.arkcdp.execute(cmd)

    async def execute_and_wait_for(
        self,
        cmd: Generator[dict, T, Any],
        event_type: Type[TEvent],
        buffer_size=10,
    ) -> Tuple[T, Type[TEvent]]:
        async with self.wait_for(event_type, buffer_size) as proxy:
            result = await self.execute(cmd)
        return (result, proxy.value)

    async def execute_and_wait_paused(
        self, cmd: Generator[dict, T, Any]
    ) -> Tuple[T, debugger.Paused]:
        return await self.execute_and_wait_for(cmd, debugger.Paused)

    def _listen(self, arkcdp: ArkCdp):
        self.arkcdp = arkcdp

        async def onExecutionContextCreated(event: runtime.ExecutionContextCreated):
            self.execution_context = event.context

        async def onScriptParsed(event: debugger.ScriptParsed):
            if event.script_id not in self.scripts:
                self.scripts[event.script_id] = event

        async def onPaused(event: debugger.Paused):
            self.paused = event

        async def onResumed(event: debugger.Resumed):
            self.paused = None

        handlers = {
            runtime.ExecutionContextCreated: onExecutionContextCreated,
            debugger.ScriptParsed: onScriptParsed,
            debugger.Paused: onPaused,
            debugger.Resumed: onResumed,
        }
        for k, v in handlers.items():
            self.arkcdp.listen(k, v)

    async def get_code(self, script_id: runtime.ScriptId) -> str:
        if script_id not in self.script_sources:
            result: str = await self.arkcdp.execute(
                debugger.get_script_source(script_id)
            )
            self.script_sources[script_id] = result
        return self.script_sources[script_id]

    async def get_remote_object(self, id: runtime.RemoteObjectId) -> Generator[
        T_JSON_DICT,
        T_JSON_DICT,
        Tuple[
            List[runtime.PropertyDescriptor],
            Optional[List[runtime.InternalPropertyDescriptor]],
            Optional[List[runtime.PrivatePropertyDescriptor]],
            Optional[runtime.ExceptionDetails],
        ],
    ]:
        return await self.arkcdp.execute(runtime.get_properties(id))

    def current_scope(self) -> Optional[debugger.Scope]:
        if self.paused is None:
            return None
        frame: debugger.CallFrame = next(self.paused.call_frames, None)
        return next((x for x in frame.scope_chain if x.type_ == "local"), None)


def _panel(
    renderable: RenderableType,
    title: Optional[TextType] = None,
    # box: box.Box = box.SIMPLE,
    height: Optional[int] = None,
) -> Panel:
    return Panel(
        title=f"[green]{title}[/green]",
        title_align="left",
        # box=box,
        renderable=renderable,
        height=height,
    )


@dataclass
class StackFrame:
    id: debugger.CallFrameId
    name: str
    url: str
    line: int

    @property
    def file(self) -> str:
        return pathlib.Path(self.url).stem

    def __repr__(self) -> str:
        return f"{self.id} {self.name} {self.file}:{self.line}"


def stackframe(call_frame: debugger.CallFrame) -> StackFrame:
    return StackFrame(
        id=call_frame.call_frame_id,
        name=call_frame.function_name,
        url=call_frame.url,
        line=call_frame.location.line_number + 1,
    )


class ArkDebuggerLayout:
    def __init__(self) -> None:
        self.code = Layout(name="code")
        self.scope = Layout(name="scope")
        self.callstack = Layout(name="callstack")

        paused = Layout(name="paused")
        paused.split_row(self.scope, self.callstack)

        root = Layout(name="root")
        root.split_row(self.code, paused)
        self.root = Panel(root, height=30, box=box.SIMPLE)

    def update_code(
        self,
        code: str,
        url: Optional[str] = None,
        breakpoint_lines: Set[int] = [],
        current_line: Optional[int] = None,
    ):
        syntax = Syntax(
            code=code,
            lexer="ts",
            line_numbers=True,
            highlight_lines=[l + 1 for l in breakpoint_lines],
        )
        if current_line is not None:
            syntax.stylize_range(
                style=Style(bgcolor="yellow4"),
                start=(current_line + 1, 0),
                end=(current_line + 2, 0),
            )
        panel = _panel(
            title=Text(url, style="repr.filename"),
            renderable=syntax,
        )
        self.code.update(panel)

    def update_callstack(
        self,
        frames: Iterable[StackFrame],
        active_id: Optional[debugger.CallFrameId] = None,
    ) -> None:
        table = Table(
            Column("#", justify="right"),
            Column("function"),
            Column("location", justify="right"),
            box=box.SIMPLE,
            show_edge=False,
        )
        for frame in frames:
            style = None
            arrow = ""
            if frame.id == active_id:
                style = "bold"
                arrow = "[red]:arrow_forward:[/red] "

            table.add_row(
                f"{arrow}{frame.id}",
                frame.name,
                f"{frame.file}:{frame.line}",
                style=style,
            )
        self.callstack.update(_panel(table, title="call stack"))

    def update_scope(
        self,
        scope_chain: Iterable[debugger.Scope],
        remoteObjects: Dict[
            runtime.RemoteObjectId,
            Iterable[runtime.PropertyDescriptor],
        ],
    ):
        table = Table(
            Column("name"),
            Column("value"),
            box=box.SIMPLE,
            show_edge=False,
        )
        for scope in scope_chain:
            table.add_row(f":red_triangle_pointed_down: {scope.type_}", style="bold")
            vars = list(remoteObjects.get(scope.object_.object_id))
            if len(vars):
                for prop in vars:
                    table.add_row(prop.name, Pretty(prop.value.value))
            else:
                table.add_row("none", "", style="italic")
            table.add_section()
        self.scope.update(_panel(table, title="variables"))


class ArkDebuggerLive:

    def __init__(self, debug: ArkDebugger) -> None:
        self.debug = debug
        self._layout = ArkDebuggerLayout()
        pass

    @property
    def layout(self) -> Layout:
        return self._layout.root

    async def update(self):
        paused = self.debug.paused
        frame: debugger.CallFrame = next(iter(paused.call_frames))
        code = await self.debug.get_code(frame.location.script_id)
        self._layout.update_code(
            code=code,
            url=frame.url,
            current_line=frame.location.line_number,
        )
        frames = [stackframe(frame) for frame in paused.call_frames]
        self._layout.update_callstack(frames, frame.call_frame_id)

        async def get_props(
            id: runtime.RemoteObjectId,
        ) -> List[runtime.PropertyDescriptor]:
            props, _, _, _ = await self.debug.get_remote_object(id)
            return props

        remote_objects = {
            scope.object_.object_id: await get_props(scope.object_.object_id)
            for scope in frame.scope_chain
        }
        self._layout.update_scope(frame.scope_chain, remote_objects)
